(* ::Package:: *)

(* ::Title:: *)
(*Tweezer Aberrations using Kant' s Formalism*)


(* ::Subtitle:: *)
(*Description*)


(* ::Text:: *)
(*Gabriel Patenotte*)
(*April 11, 2024*)
(*Ni Lab*)
(*Minor adjustments by David Wellnitz*)
(**)
(*Version: *)
(*2 (4/19). Speeding up numerical integration, changing definition of \[CapitalPhi], gradient descent to find max energy with aberrations*)
(**)
(*References:*)
(*1.Kant, R. An Analytical Solution of Vector Diffraction for Focusing Optical Systems with Seidel Aberrations. J. Mod. Opt. 40, 2293\[Dash]2310 (1993).*)
(*2.Kant, R. An Analytical Method of Vector Diffraction for Focusing Optical Systems with Seidel Aberrations II: Astigmatism and Coma. J. Mod. Opt. 42, 299\[Dash]320 (1995).*)
(*3.Singh, R. K., Senthilkumaran, P. & Singh, K. Tight focusing of vortex beams in presence of primary astigmatism. J. Opt. Soc. Am. A 26, 576 (2009).*)
(*4. Richards, B. & Wolf, E. Electromagnetic diffraction in optical systems, II. Structure of the image field in an aplanatic system. (1959).*)
(*5. Wolf, E. Electromagnetic diffraction in optical systems - I. An integral representation of the image field. Proc. R. Soc. Lond. Ser. A Math. Phys. Sci. 253, 349\[Dash]357 (1959).*)
(*    *)
(*Reference 3 is based on the other references and is most applicable to this problem, 1 and 2 expand 4 and 5 in the presence of aberrations, 4 and 5 provide the most general answer.*)


(* ::Chapter:: *)
(*Constants*)


(* ::Subsection:: *)
(*Units*)


(* ::Text:: *)
(*Defining length to be in units of nm, time in units of \[Mu]s, mass in units of amu, and charge in units of electron charge.*)


(* ::Input::Initialization:: *)
length=Quantity["Nanometers"];
mass = Quantity["AtomicMassUnit"];
e = charge=Quantity["ElementaryCharge"];
time = Quantity["Microseconds"];
frequency=1/time;
energy=(mass length^2)/time^2;


(* ::Input::Initialization:: *)
nm=Quantity["Nanometers"]/length;
\[Mu]m=um=Quantity["Micrometers"]/length;
mm=Quantity["Millimeters"]/length;
cm=Quantity["Centimeters"]/length;
amu=Quantity["AtomicMassUnit"]/mass;
\[Mu]s=Quantity["Microseconds"]/time;
MHz=Quantity["Megahertz"]/frequency;
GHz=Quantity["Gigahertz"]/frequency;
kHz=Quantity["Kilohertz"]/frequency;
eV=Quantity["Electronvolts"]/energy;
icm=Quantity[("Centimeters")^-1]length;
deg=\[Pi]/180;(*to express angles in degrees but have them be in radians*)
volt=Quantity["Volts"] (charge time^2)/(mass length^2);
mV=Quantity["Millivolts"] (charge time^2)/(mass length^2);
MV=Quantity["Megavolts"] (charge time^2)/(mass length^2);
hartree=Quantity["HartreeEnergy"]/energy;
auPol=e^2/charge^2 a0^2/hartree;
MW=Quantity["Megawatts"] time/energy;
kW=Quantity["Kilowatts"] time/energy;
debye=Quantity["Debyes"]/(charge length)//N;


(* ::Subsection:: *)
(*Fundamental constants*)


(* ::Input::Initialization:: *)
c=Quantity["SpeedOfLight"] time/length;
h=Quantity["PlanckConstant"] time/(mass length^2);
a0=Quantity["Bohr radius"]/length;
\[Epsilon]0=Quantity["Permittivity of free space"] (mass length^3)/(time^2 charge^2);
\[HBar]=h/2\[Pi];


(* ::Subsection:: *)
(*Experimental parameters*)


(* ::Input::Initialization:: *)
\[Chi]m=1/2 ArcCos[1/3]; (*Magic ellipticity*)
\[Alpha]par=1872.12 auPol;
\[Alpha]perp=467.038 auPol;
U=4 h MHz;(*trap depth [energy] *)
\[Lambda]=1064 nm; (*Tweezer wavelength*)
k=2\[Pi]/\[Lambda];
coordinates = {ivec->{1,0,0},jvec->{0,1,0},kvec->{0,0,1}};


(* ::Chapter:: *)
(*Simulation Parameters*)


(* ::Text:: *)
(*Aberrations order: Astigmatism, Coma, spherical, field curvature, distortion*)


(* ::Input::Initialization:: *)
ChiExp=N[1/2 ArcCos[1/3]];
aberrations={0.1,0,0,0,0}\[Lambda];
GridStep=50.nm;
GridMaxPerp=400.nm;
exp=Association[f->18 mm, \[Alpha]->ArcSin[0.55],A->Sin[ChiExp],B->Cos[ChiExp]I];


(* ::Chapter:: *)
(*Functions*)


(* ::Text:: *)
(*All formulas come from the cited references. Conceptually, they define an aberration free wavefront that is f [length] away from the origin, and an actual wavefront r [length] that is warped by some amount \[CapitalPhi] [length]. The idea is that a plane wave emits from every point on r, which creates a maximum intensity at the origin. The warping \[CapitalPhi] will be different depending on the type of aberration. For instance if the only aberration is coma, set uc >0 and other amplitudes (u) to zero.*)


(* ::Text:: *)
(*Aberration amplitudes in units of length: ua: astigmatism; uc: coma; us: spherical; uf: curvature of field; ud: distortion. These should typically be fractions of the laser wavelength.*)


(* ::Input::Initialization:: *)
\[Rho][\[Theta]_,\[Alpha]_]=Sin[\[Theta]]/Sin[\[Alpha]] (* Zonal radius, 2.12 *);
\[CapitalPhi][\[Theta]_,\[Phi]_,\[Alpha]_,{ua_,uc_,us_,uf_,ud_}]=ua \[Rho][\[Theta],\[Alpha]]^2 Cos[\[Phi]]^2+uc \[Rho][\[Theta],\[Alpha]]^3 Cos[\[Phi]]+us \[Rho][\[Theta],\[Alpha]]^4+uf \[Rho][\[Theta],\[Alpha]]^2+ud \[Rho][\[Theta],\[Alpha]] Cos[\[Phi]] ;
r[\[Theta]_,\[Phi]_,f_,\[Alpha]_,{ua_,uc_,us_,uf_,ud_}]=f+\[CapitalPhi][\[Theta],\[Phi],\[Alpha],{ua,uc,us,uf,ud}] (* Wavefront position *);
\[CapitalTheta][\[Theta]_,\[Phi]_,f_,\[Alpha]_,{ua_,uc_,us_,uf_,ud_}]=1/(2r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]^2) (D[r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}],\[Theta]]^2+D[r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}],\[Phi]]^2/Sin[\[Theta]]^2) ;(* Normalization Theta, 2.22 *);
\[Sigma][\[Theta]_,\[Phi]_,f_,\[Alpha]_,{ua_,uc_,us_,uf_,ud_}]=1/(1-\[CapitalTheta][\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]) (* Normalization sigma, 2.21 *);
sx[\[Theta]_,\[Phi]_,f_,\[Alpha]_,{ua_,uc_,us_,uf_,ud_}]=1/\[Sigma][\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}] (Sin[\[Theta]]Cos[\[Phi]]-D[r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}],\[Theta]]/r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}] Cos[\[Theta]]Cos[\[Phi]]+D[r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}],\[Phi]]/(r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]Sin[\[Theta]]) Sin[\[Phi]]) (* x-component of vector normal to wavefront, 2.17 *);
sy[\[Theta]_,\[Phi]_,f_,\[Alpha]_,{ua_,uc_,us_,uf_,ud_}]=1/\[Sigma][\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}] (Sin[\[Theta]]Sin[\[Phi]]-D[r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}],\[Theta]]/r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}] Cos[\[Theta]]Sin[\[Phi]]-D[r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}],\[Phi]]/(r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]Sin[\[Theta]]) Cos[\[Phi]])  (* y-component of vector normal to wavefront, 2.17 *);
sz[\[Theta]_,\[Phi]_,f_,\[Alpha]_,{ua_,uc_,us_,uf_,ud_}]=1/\[Sigma][\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}] (Cos[\[Theta]]+D[r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}],\[Theta]]/r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}] Sin[\[Theta]])  (* z-component of vector normal to wavefront, 2.17 *);
a[\[Theta]_,\[Phi]_,f_,\[Alpha]_,{ua_,uc_,us_,uf_,ud_}]=f l0 Sqrt[sz[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]](((sx[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]^2 sz[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]+sy[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]^2)ivec+sx[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]sy[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}](sz[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]-1)jvec)/(sx[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]^2+sy[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]^2)-sx[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]kvec)/.coordinates (* E-field strength vector, 2.24 *);
fx[\[Theta]_,\[Phi]_,f_,\[Alpha]_,{ua_,uc_,us_,uf_,ud_}]=(-(D[r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}],\[Theta]]/r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}])Cos[\[Theta]]Cos[\[Phi]]+D[r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}],\[Phi]]/(r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]Sin[\[Theta]]) Sin[\[Phi]])(* x-component of wavefront deviation from normal, 3.6b *);
fy[\[Theta]_,\[Phi]_,f_,\[Alpha]_,{ua_,uc_,us_,uf_,ud_}]=(-(D[r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}],\[Theta]]/r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}])Cos[\[Theta]]Sin[\[Phi]]-D[r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}],\[Phi]]/(r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]Sin[\[Theta]]) Cos[\[Phi]])(* y-component of wavefront deviation from normal, 3.6b *);
fz[\[Theta]_,\[Phi]_,f_,\[Alpha]_,{ua_,uc_,us_,uf_,ud_}]=(D[r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}],\[Theta]]/(r[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]Sin[\[Theta]]) Sin[\[Theta]])(* z-component of wavefront deviation from normal, 3.6b *);
fvec[\[Theta]_,\[Phi]_,f_,\[Alpha]_,{ua_,uc_,us_,uf_,ud_}]={fx[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}],fy[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}],fz[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]};
A2[\[Theta]_,\[Phi]_,f_,\[Alpha]_,{ua_,uc_,us_,uf_,ud_}]=Sqrt[sz[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]] (* apodization factor, 3.5b *);
P[\[Theta]_,\[Phi]_,f_,\[Alpha]_,{ua_,uc_,us_,uf_,ud_},A_,B_]=1/(sx[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]^2+sy[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]^2) ({
 {A(sy[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]^2+sx[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]^2 sz[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}])+B(-sx[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]sy[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]+sx[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]sy[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]sz[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}])},
 {A(-sx[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]sy[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]+sx[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]sy[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]sz[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}])+B(sx[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]^2+sy[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]^2 sz[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}])},
 {(A(-sx[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}])+B(-sy[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]))(sx[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]^2+sy[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]^2)}
})(* polarization distribution of input beam at exit pupil plane, 3.5b *);
ap[\[Theta]_,\[Phi]_,f_,\[Alpha]_,{ua_,uc_,us_,uf_,ud_},A_,B_]=A2[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]P[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud},A,B];
n[\[Theta]_,\[Phi]_]=Sin[\[Theta]]Cos[\[Phi]]ivec+Sin[\[Theta]]Sin[\[Phi]]jvec+Cos[\[Theta]]kvec/.coordinates(*normal vector to non-aberrated/Gaussian wavefront*);
rP[rp_,\[Theta]p_,\[Phi]p_]=rp(Sin[\[Theta]p]Cos[\[Phi]p]ivec+Sin[\[Theta]p]Sin[\[Phi]p]jvec+Cos[\[Theta]p]kvec)/.coordinates (* position vector relative to image plane origin *);
v[rp_,\[Theta]p_,\[Alpha]_]=k rp Sin[\[Theta]p]Sin[\[Alpha]];
u[rp_,\[Theta]p_,\[Alpha]_]=k rp Cos[\[Theta]p]Sin[\[Alpha]]^2;
jacobian[\[Theta]_,\[Phi]_,f_,\[Alpha]_,{ua_,uc_,us_,uf_,ud_}]=(D[sx[\[Theta]0,\[Phi]0,f,\[Alpha],{ua,uc,us,uf,ud}],\[Theta]0]D[sy[\[Theta]0,\[Phi]0,f,\[Alpha],{ua,uc,us,uf,ud}],\[Phi]0]-D[sx[\[Theta]0,\[Phi]0,f,\[Alpha],{ua,uc,us,uf,ud}],\[Phi]0]D[sy[\[Theta]0,\[Phi]0,f,\[Alpha],{ua,uc,us,uf,ud}],\[Theta]0])/.{\[Theta]0->\[Theta],\[Phi]0->\[Phi]};
eIntegrand[\[Theta]_,\[Phi]_,rp_,\[Theta]p_,\[Phi]p_,f_,\[Alpha]_,{ua_,uc_,us_,uf_,ud_},A_,B_]=Flatten[A2[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]P[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud},A,B]Exp[I k (\[CapitalPhi][\[Theta],\[Phi],\[Alpha],{ua,uc,us,uf,ud}]+((\[CapitalTheta][\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]n[\[Theta],\[Phi]]-1/\[Sigma][\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}] fvec[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]) . rP[rp,\[Theta]p,\[Phi]p]))]Exp[-I (v[rp,\[Theta]p,\[Alpha]] Sin[\[Theta]])/Sin[\[Alpha]] Cos[\[Phi]-\[Phi]p]]Exp[-I (u[rp,\[Theta]p,\[Alpha]]Cos[\[Theta]])/Sin[\[Alpha]]^2]jacobian[\[Theta],\[Phi],f,\[Alpha],{ua,uc,us,uf,ud}]];


(* ::Chapter:: *)
(*Analyze E-field*)


(* ::Subsection:: *)
(*Energy of rotational states*)


(* ::Text:: *)
(*Polarizability of NaCs along parallel and perpendicular directions. \[Alpha]0 is the scalar polarizability, \[Alpha]2 is the tensor polarizability. The other equations are the probabilities of the rotational eigenstates.*)
(**)
(*The eigenstates of N=1 are labeled as lower, middle, and upper. Lower is aligned along the major axis of polarization and is therefore a 'low' field seeking state.*)


(* ::Input::Initialization:: *)
\[CapitalDelta]=\[Alpha]par-\[Alpha]perp;
(*\[Alpha]0=(2\[Alpha]perp+\[Alpha]par)/3*);
\[Alpha]0=1074.8 auPol;
(*\[Alpha]2=2/3 \[CapitalDelta]*);
\[Alpha]2=1800 auPol;
\[Alpha]N0[\[Chi]_]=\[Alpha]0;
\[Alpha]N1U[\[Chi]_]=\[Alpha]0-\[Alpha]2 1/5;
\[Alpha]N1L[\[Chi]_]=\[Alpha]0+\[Alpha]2 (1+3Cos[2\[Chi]])/10;
\[Alpha]N1M[\[Chi]_]=\[Alpha]0-\[Alpha]2 (-1+3Cos[2\[Chi]])/10;
I0=(2\[Epsilon]0 c U)/\[Alpha]0;
I0 ((cm^2)/kW)(*Intensity of tweezer in kW/cm^2 to produce a scalar depth of U*)

fnString="dr"<>ToString[GridStep/nm]<>"-rmax"<>ToString[GridMaxPerp/nm]<>"-astig"<>ToString[aberrations[[1]]]<>"-dchi"<>ToString[ChiExp-\[Chi]m]<>"a0"<>ToString[\[Alpha]0/auPol]<>"a2"<>ToString[\[Alpha]2/auPol]



(* ::Input::Initialization:: *)
UN0[int_,\[Chi]_]=\[Alpha]N0[\[Chi]]/(2\[Epsilon]0 c) int;
UN1L[int_,\[Chi]_]=\[Alpha]N1L[\[Chi]]/(2\[Epsilon]0 c) int;
UN1M[int_,\[Chi]_]=\[Alpha]N1M[\[Chi]]/(2\[Epsilon]0 c) int;
UN1U[int_,\[Chi]_]=\[Alpha]N1U[\[Chi]]/(2\[Epsilon]0 c) int;


(* ::Input::Initialization:: *)
SetDirectory[NotebookDirectory[]];


(* ::Input::Initialization:: *)
Needs["Polarization`"]
Needs["Jones`"]


(* ::Input::Initialization:: *)
FindEnergy[eField_,\[Alpha]State_]:=Module[
{energies,eFieldNorm,intensity,\[Chi]},
eFieldNorm=Normalize[eField];
intensity=eField . eField\[Conjugate];
\[Chi]=FindEllipticity[eFieldNorm];
energies=Re[intensity \[Alpha]State[\[Chi]]/\[Alpha]0 ]]


(* ::Subsection:: *)
(*Find fastest numerical integration method*)


(* ::Text:: *)
(*Only run this section to find the fastest integration method. Plug in values that are similiar to those that will be used throughout this notebook.*)


(* ::Text:: *)
(*Note: The integral is a Levin integral so NIntegrate will always try to use a LevinRule method. However, it fails due to the large number of sinusoidal functions. I tried LevinFunctions of ExpRelated and BesselRelated which works, but is slow. When LevinRule fails it uses MultiDimensionalRule which is fast. However, I found that the other methods: GaussKronrod, GaussBerntsen, and especially ClenshawCurtis, are faster when the presionGoal/accuracyGoal exceed 4. MultiDimensional is best is the precision goal is below 4. GaussKronrod is best when the precision goal is 4.*)


(* ::Input::Initialization:: *)
FindBestMethod[f_,\[Alpha]_,goal_,aberrations_,{rp_,\[Theta]p_,\[Phi]p_},A_,B_,methods_]:=Module[
(*f: objective focal length; \[Alpha]: objective NA; A/B: Input laser polarization*)
{integrand,\[Theta],\[Phi],timingsByCoordinate,timingsAllCoordinates,bestMethod},
integrand[\[Theta]_,\[Phi]_]=eIntegrand[\[Theta],\[Phi],rp,\[Theta]p,\[Phi]p,f,\[Alpha],aberrations,A,B];
timingsByCoordinate=Table[(Timing[NIntegrate[integrand[\[Theta]0,\[Phi]0][[p]],{\[Theta]0,0,\[Alpha]},{\[Phi]0,0,2\[Pi]},Method->#,PrecisionGoal->goal,AccuracyGoal->goal,WorkingPrecision->goal]]&/@methods)[[All,1]]
,{p,1,3}];
timingsAllCoordinates=Total[timingsByCoordinate];
bestMethod=methods[[Position[timingsAllCoordinates,Min[timingsAllCoordinates]][[1,1]]]]
]


(* ::Input::Initialization:: *)
methods={{"MultidimensionalRule","SymbolicProcessing"->1},{"ClenshawCurtisRule","Points"->15},{"LobattoKronrodRule","Points"->9},{"GaussKronrodRule","Points"->10},{"GaussBerntsenEspelidRule","Points"->9}};
goal=8;
bestMethod=FindBestMethod[18 mm, 0.55, goal,aberrations, {1578 nm,Pi,0},1,0,methods]


(* ::Subsection:: *)
(*Find e-field*)


(* ::Text:: *)
(*Cartesian to spherical coordinate transformation.*)


(* ::Input::Initialization:: *)
sgn[x_]:=If[x>=0,1,-1];
rp[xp_,yp_,zp_]:=Sqrt[xp^2+yp^2+zp^2];
\[Theta]p[xp_,yp_,zp_]:=If[xp^2+yp^2+zp^2==0,0,ArcCos[zp/Sqrt[xp^2+yp^2+zp^2]]];
\[Phi]p[xp_,yp_,zp_]:=If[xp^2+yp^2==0,0,sgn[yp]ArcCos[xp/Sqrt[xp^2+yp^2]]];


(* ::Input::Initialization:: *)
eMax=With[{xp=0,yp=0,zp=0},NIntegrate[eIntegrand[\[Theta]0,\[Phi]0,rp[xp,yp,zp],\[Theta]p[xp,yp,zp],\[Phi]p[xp,yp,zp],exp[f],exp[\[Alpha]],{0,0,0,0,0},exp[A],exp[B]],{\[Theta]0,0,exp[\[Alpha]]},{\[Phi]0,0,2\[Pi]},Method->bestMethod,PrecisionGoal->goal,AccuracyGoal->goal,WorkingPrecision->goal]];
iMax=eMax . eMax\[Conjugate];
FindField[{xp_,yp_,zp_},exp_,aberrations_,method_]:=Module[
{integrand,\[Theta],\[Phi]},
integrand[\[Theta]_,\[Phi]_]=1/Sqrt[iMax] eIntegrand[\[Theta],\[Phi],rp[xp,yp,zp],\[Theta]p[xp,yp,zp],\[Phi]p[xp,yp,zp],exp[f],exp[\[Alpha]],aberrations,exp[A],exp[B]];
NIntegrate[integrand[\[Theta]0,\[Phi]0],{\[Theta]0,0,exp[\[Alpha]]},{\[Phi]0,0,2\[Pi]},Method->method,PrecisionGoal->goal,AccuracyGoal->goal,WorkingPrecision->goal]
]


(* ::Text:: *)
(*Find point of maximum intensity for the given aberrations*)


(* ::Input::Initialization:: *)
dxs[\[Gamma]_]:=\[Gamma] DiagonalMatrix[{1,1,1}];
FindGradient[x0_,\[Alpha]State_,exp_,aberrations_,method_]:=Module[
{xs,x1s,es,ens},
x1s=(x0+dxs[1 nm][[#]])&/@{1,2,3};
xs=Join[{x0},x1s];
es=Chop[FindField[#,exp,aberrations,method]&/@xs,10^-goal];
ens=FindEnergy[#,\[Alpha]State]&/@es;
Table[(ens[[i]]-ens[[1]])/(1 nm),{i,2,4}]
];
Find\[Gamma]New[xNew_,xOld_,dENew_,dEOld_]:=Abs[(xNew-xOld) . (dENew-dEOld)]/(dENew-dEOld) . (dENew-dEOld);
FindxNew[xOld_,\[Gamma]Old_,dEOld_]:=xOld+\[Gamma]Old dEOld;


(* ::Text:: *)
(*Gradient descent to find point of maximum energy. xs are the positions, \[Gamma]s are the learning rate, and dEs are the gradients. *)


(* ::Input::Initialization:: *)
FindPeak[x0_,exp_,aberrations_,\[Alpha]State_,method_]:=Module[
{xOld,dEOld,\[Gamma]Old,xs,\[Gamma]New,dENew,xNew,xPeak,fieldPeak,enPeak,lp,energyNearPeak,\[Omega]Peak},
	xOld=x0;
	dEOld=FindGradient[xOld,\[Alpha]State,exp,aberrations,method];
	\[Gamma]Old=(10 nm)/Norm[dEOld];
	xs={xOld};
	While[
	Norm[dEOld]>10^-8,
		xNew=FindxNew[xOld,\[Gamma]Old,dEOld];
		dENew=FindGradient[xNew,\[Alpha]State,exp,aberrations,method];
		\[Gamma]New=Find\[Gamma]New[xNew,xOld,dENew,dEOld];
		\[Gamma]Old=\[Gamma]New;
		xOld=xNew;
		dEOld=dENew;
		AppendTo[xs,xOld];
	];
lp=ListPlot[xs\[Transpose],Joined->True,PlotRange->All,PlotRange->All,PlotLegends->{"x","y","z"},Frame->True,BaseStyle->{FontSize->16,FontColor->Black},FrameLabel->{"Step","dx"},PlotLabel->"Convergence to point\n of minimum energy"];
xPeak = xs[[-1]];
fieldPeak=FindField[xPeak,exp,aberrations,method];
enPeak=FindEnergy[fieldPeak,\[Alpha]State];
energyNearPeak=FindEnergy[FindField[xPeak+(50 nm DiagonalMatrix[{1,1,1}])[[#]],exp,aberrations,method],\[Alpha]State]&/@{1,2,3};
\[Omega]Peak=Sqrt[(2Abs[enPeak-energyNearPeak]/(156amu))/(50 nm)^2];
{xPeak,fieldPeak,enPeak,\[Omega]Peak,lp}

]


(* ::Input::Initialization:: *)
{xPeak0,fieldPeak0,enPeak0,\[Omega]Peak0,peakPlot0}=FindPeak[{-0.,-0.,-0.}nm,exp,aberrations,\[Alpha]N0,bestMethod];


(* ::Input::Initialization:: *)
bestMethod=FindBestMethod[18 mm, 0.55, goal,aberrations, {Abs[xPeak0[[3]]],Pi/2*(1-Sign[xPeak0[[3]]]),0},1,0,methods]


(* ::Subsubsection:: *)
(*Potential around the peak*)


(* ::Input::Initialization:: *)
GridMaxZTmp=Sqrt[Sqrt[\[Omega]Peak0[[1]] \[Omega]Peak0[[2]]]/\[Omega]Peak0[[3]]] GridMaxPerp;
GridMaxZ=Floor[GridMaxZTmp/GridStep] GridStep;


(* ::Text:: *)
(*FindEnergy returns the potential in arbitrary units. Experimentally, the trap depth is 4MHz. Thus, we rescale by enPeak0 to get 4MHz. This is the potential in Frequency units (MHz)*)


(* ::Input::Initialization:: *)
EnergyGridN0=Table[{x,y,z,FindEnergy[FindField[xPeak0+{x,y,z},exp,aberrations,bestMethod],\[Alpha]N0]/enPeak0*U/h},{x,-GridMaxPerp,GridMaxPerp,GridStep},{y,-GridMaxPerp,GridMaxPerp,GridStep},{z,-GridMaxZ,GridMaxZ,GridStep}];
of="/path/to/dir/";
Export[of<>"N0-"<>fnString<>".h5",EnergyGridN0];


(* ::Input::Initialization:: *)
EnergyGridN1M=Table[{x,y,z,FindEnergy[FindField[xPeak0+{x,y,z},exp,aberrations,bestMethod],\[Alpha]N1M]/enPeak0*U/h},{x,-GridMaxPerp,GridMaxPerp,GridStep},{y,-GridMaxPerp,GridMaxPerp,GridStep},{z,-GridMaxZ,GridMaxZ,GridStep}];
Export[of<>"N1M-"<>fnString<>".h5",EnergyGridN1M];
